#!/bin/bash
# DeepGP 安装脚本（项目根目录）
# 自动选择国内或国际版本

echo "=========================================="
echo "  DeepGP 环境管理"
echo "=========================================="
echo ""
echo "请选择操作:"
echo "  1) 全新安装（国内用户，推荐，CPU 版本 - 稳定）"
echo "  2) 全新安装（国际用户，CPU 版本）"
echo "  3) 快速修复当前环境（安装缺失依赖 + CPU 版本）"
echo "  4) 完全重装（删除旧环境后重新安装 CPU 版本）"
echo ""
read -p "请输入选项 (1/2/3/4): " choice

case $choice in
    1)
        echo "使用国内优化版本安装..."
        bash scripts/setup_env_china.sh
        ;;
    2)
        echo "使用国际版本安装..."
        bash scripts/setup_env.sh
        ;;
    3)
        echo "快速修复当前环境（CPU 版本）..."
        eval "$(conda shell.bash hook)"
        conda activate deepgp
        if [ $? -ne 0 ]; then
            echo "❌ 无法激活 deepgp 环境，请先运行选项 1 或 2"
            exit 1
        fi
        
        echo ""
        echo "============================================"
        echo "  先安装兼容的 C++ 标准库（重要！）"
        echo "============================================"
        conda install -c conda-forge libstdcxx-ng -y
        echo "✓ C++ 标准库安装完成"
        
        echo ""
        echo "安装缺失的依赖..."
        pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
        pip install -r streamlit_app/requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
        
        echo ""
        echo "重新安装 scipy 和 scikit-learn（确保兼容性）..."
        pip install scipy scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple --force-reinstall --no-deps
        pip install scipy scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple --no-deps
        
        # ⚠️ 关键：最后再次确保 pandas 和 numpy 版本正确（防止被覆盖）
        echo ""
        echo "最终确保 pandas 和 numpy 版本正确..."
        pip install "numpy<2.0" "pandas<2.0" -i https://pypi.tuna.tsinghua.edu.cn/simple --force-reinstall
        
        echo ""
        echo "============================================"
        echo "  卸载旧版本 DGL（完全清理）"
        echo "============================================"
        
        pip uninstall dgl -y 2>/dev/null
        conda uninstall dgl -y 2>/dev/null
        pip cache purge 2>/dev/null
        conda clean --all -y 2>/dev/null
        
        # 删除可能残留的 DGL 文件
        rm -rf /opt/miniconda3/envs/deepgp/lib/python3.11/site-packages/dgl* 2>/dev/null
        
        echo ""
        echo "============================================"
        echo "  安装 DGL CPU 版本（conda-forge 可靠源）"
        echo "============================================"
        
        # 方法1：使用 conda-forge（CPU 版本，无 CUDA 依赖）
        echo "尝试 conda-forge 安装（推荐）..."
        conda install -c conda-forge dgl -y
        
        if [ $? -ne 0 ]; then
            echo ""
            echo "⚠️  conda-forge 失败，尝试官方 PyPI..."
            # 方法2：官方 PyPI（明确 CPU 版本）
            pip install dgl --no-cache-dir
        fi
        
        echo ""
        echo "============================================"
        echo "  修复依赖版本（兼容性）"
        echo "============================================"
        
        # 重要：降级 NumPy 和 pandas（兼容 PyTorch 2.1.2）
        pip install "numpy<2.0" "pandas<2.0" -i https://pypi.tuna.tsinghua.edu.cn/simple --force-reinstall
        
        echo ""
        echo "============================================"
        echo "  验证安装"
        echo "============================================"
        
        python -c "
import dgl
import torch
import numpy
print(f'✓ DGL: {dgl.__version__}')
print(f'✓ PyTorch: {torch.__version__}')
print(f'✓ NumPy: {numpy.__version__}')
print(f'✓ DGL Backend: {dgl.backend.backend_name}')
try:
    import fastNLP
    print(f'✓ FastNLP: OK')
except:
    print(f'⚠️  FastNLP: 导入失败')
" 2>/dev/null || echo "⚠️  验证失败，请检查错误信息"
        
        echo "✓ 修复完成！"
        ;;
    4)
        echo "⚠️  这将删除旧的 deepgp 环境并重新安装"
        read -p "确认继续? (y/n): " confirm
        if [ "$confirm" = "y" ]; then
            conda remove -n deepgp --all -y
            echo "重新安装..."
            bash scripts/setup_env_china.sh
        else
            echo "取消操作"
        fi
        ;;
    *)
        echo "无效选项，默认使用国内优化版本..."
        bash scripts/setup_env_china.sh
        ;;
esac

